import pandas as pd
import math
import matplotlib.pyplot as plt
import numpy as np
from sklearn.metrics import r2_score


# 定义填补函数
def fill_sediment(row, df):
    if pd.isna(row['含沙量(kg/m3) ']):
        # 获取当前行的位置
        index = df.index.tolist().index(row.name)
        d_forward = 1
        d_backward = 1
        # 获取 DataFrame 的最大位置
        max_index = len(df) - 1
        min_index = 0

        # 查找下一个非空值
        next_non_na_index = None
        while index + d_forward <= max_index:
            if not pd.isna(df.iloc[index + d_forward]['含沙量(kg/m3) ']):
                next_non_na_index = index + d_forward
                break
            d_forward += 1

        # 查找前一个非空值
        prev_non_na_index = None
        while index - d_backward >= min_index:
            if not pd.isna(df.iloc[index - d_backward]['含沙量(kg/m3) ']):
                prev_non_na_index = index - d_backward
                break
            d_backward += 1

        # 如果找到了前后非空值
        if next_non_na_index is not None and prev_non_na_index is not None:
            q_st_prev = df.iloc[prev_non_na_index]['含沙量(kg/m3) ']
            q_st_next = df.iloc[next_non_na_index]['含沙量(kg/m3) ']
            return (d_forward / (d_forward + d_backward)) * q_st_prev + (d_backward / (d_forward + d_backward)) * q_st_next
        # 如果只找到了前一个非空值
        elif prev_non_na_index is not None:
            return df.iloc[prev_non_na_index]['含沙量(kg/m3) ']
        # 如果只找到了下一个非空值
        elif next_non_na_index is not None:
            return df.iloc[next_non_na_index]['含沙量(kg/m3) ']

    return row['含沙量(kg/m3) ']


def data_clean():
    target_years = [2017, 2018, 2019, 2021]
    # 读取 Excel 文件
    file_path = 'ori_message.xlsx'
    try:
        df_all = pd.read_excel(file_path, sheet_name=None)
    except FileNotFoundError:
        print(f"错误：文件 {file_path} 未找到。")
    except Exception as e:
        print(f"错误：读取文件时发生未知错误：{e}")
    else:
        all_6_year_data = []
        all_water_level_sediment = []  # 用于存储水位和含沙量数据
        for year in df_all:
            print(year)
            df = df_all[year]

            # 填充年、月、日列
            df[['年', '月', '日']] = df[['年', '月', '日']].ffill()

            # 筛选每天 8:00 的数据，并创建副本
            df_8am = df[df['时间'] == '8:00'].copy()
            # 将年、月、日列转换为整数类型
            df_8am[['年', '月', '日']] = df_8am[['年', '月', '日']].astype(int)
            # 删除时间列
            df_8am.drop(['时间'], axis=1, inplace=True)

            all_dates = []
            for month in range(1, 13):
                last_day = pd.Timestamp(f'{year}-{month}-01').days_in_month
                for day in range(1, last_day + 1):
                    if not (month == 2 and day == 29):
                        all_dates.append((year, month, day))
            all_dates_df = pd.DataFrame(all_dates, columns=['年', '月', '日'])
            # 确保 all_dates_df 中的年、月、日列也是整数类型
            all_dates_df[['年', '月', '日']] = all_dates_df[['年', '月', '日']].astype(int)
            # 合并数据，缺失数据用空值填充
            df_merged = pd.merge(all_dates_df, df_8am, on=['年', '月', '日'], how='left')
            df_merged = df_merged[~((df_merged['月'] == 2) & (df_merged['日'] == 29))]
            # 填补含沙量缺失值
            df_merged['含沙量(kg/m3) '] = df_merged.apply(lambda d: fill_sediment(d, df_merged), axis=1)
            # 设置显示选项以打印所有行
            pd.set_option('display.max_rows', None)
            print(df_merged)
            all_6_year_data.append(df_merged)

            # 提取水位(m)和含沙量(kg/m3)列
            if '水位(m)' in df_merged.columns and '含沙量(kg/m3) ' in df_merged.columns:
                water_level_sediment = df_merged[['水位(m)', '含沙量(kg/m3) ']]
                all_water_level_sediment.append(water_level_sediment)

        # 合并所有年份的水位和含沙量数据
        combined_water_level_sediment = pd.concat(all_water_level_sediment, ignore_index=True)

        # 删除包含缺失值的行
        combined_water_level_sediment = combined_water_level_sediment.dropna()

        # 计算 k 值
        N = len(combined_water_level_sediment)
        k = math.ceil(1 + math.log(N) / math.log(2))
        print(f"k 的值为: {k}")

        # 计算水位组距
        water_level_min = combined_water_level_sediment['水位(m)'].min()
        water_level_max = combined_water_level_sediment['水位(m)'].max()
        water_level_class_interval = (water_level_max - water_level_min) / 13
        print(f"水位组距: {water_level_class_interval}")

        # 按顺序将水位分成 k 组
        combined_water_level_sediment['水位分组'] = pd.cut(combined_water_level_sediment['水位(m)'], bins=k)

        # 取出水位分组列并排序
        water_level_group_df = combined_water_level_sediment[['水位分组']].sort_values(by='水位分组')

        # 去除重复的水位分组
        unique_water_level_group_df = water_level_group_df.drop_duplicates()

        # 计算每组的平均含沙量
        group_avg_sediment = combined_water_level_sediment.groupby('水位分组')['含沙量(kg/m3) '].mean()

        # 将平均含沙量添加到 unique_water_level_group_df 中
        unique_water_level_group_df['平均含沙量(kg/m3)'] = unique_water_level_group_df['水位分组'].map(group_avg_sediment)

        print(unique_water_level_group_df)

        # 绘制折线图
        plt.figure(figsize=(10, 6))
        plt.plot(unique_water_level_group_df['水位分组'].astype(str), unique_water_level_group_df['平均含沙量(kg/m3)'], label='原始数据')
        plt.xlabel('水位分组')
        plt.ylabel('平均含沙量(kg/m3)')
        plt.title('水位分组与平均含沙量的关系')
        plt.xticks(rotation=45)

        # 分段拟合
        x = np.arange(len(unique_water_level_group_df))
        y = unique_water_level_group_df['平均含沙量(kg/m3)']

        # 第一组数据（第一到八个点）
        x1 = x[:8]
        y1 = y[:8]
        m = 3  # 可按需修改 m 的值
        p1 = np.polyfit(x1, y1, m)
        y_fit1 = np.polyval(p1, x1)

        # 第二组数据（第八到第十三个点）
        x2 = x[7:13]
        y2 = y[7:13]
        p2 = np.polyfit(x2, y2, m)
        y_fit2 = np.polyval(p2, x2)

        #绘制拟合曲线
        plt.plot(unique_water_level_group_df['水位分组'].astype(str)[:8], y_fit1, label=f'第一组拟合曲线 (m={m})', color='red')
        plt.plot(unique_water_level_group_df['水位分组'].astype(str)[7:13], y_fit2, label=f'第二组拟合曲线 (m={m})', color='green')

        plt.legend()
        plt.tight_layout()
        plt.show()

        # 计算 R2
        y_fit = np.concatenate([y_fit1[:-1], y_fit2])
        r21= r2_score(y[:8], y_fit1)
        r22 = r2_score(y[7:], y_fit2)


        print(f"拟合函数的 R21 值为: {r21}")
        print(f"拟合函数的 R22 值为: {r22}")

        return unique_water_level_group_df


data_clean()